import torch


class Config:
    def __init__(self):
        self.train_path = 'data/train.txt'
        self.test_path = 'data/test.txt'
        self.dev_path = 'data/dev.txt'
        self.stop_words_path = 'data/stopwords.txt'
        self.class_path = 'data/class.txt'

        self.rf_config = RfConfig()
        self.fasttext_config = FasttextConfig()
        self.bert_config = BertConfig()
        self.llm_config = LlmConfig()


class LlmConfig:
    def __init__(self):
        self.system_content = """
你是一名新闻分类审核员，任务是将新闻标题分类到以下类别之一：
finance, realty, stocks, education, science, society, politics, sports, game, entertainment。
请根据标题内容和以下关键词与示例，匹配最相关的类别。如果标题涉及教育机构但核心是社会贡献，优先归为 society。
返回 JSON 格式：{"category": "类别", "reason": "分类原因"}

类别关键词与示例：
- finance: 银行、信用卡、贷款、利率 (例: "各银行信用卡挂失费迥异")
- realty: 房产、地价、楼盘 (例: "东5环海棠公社230平准现房")
- stocks: 股市、股指、期货 (例: "金证顾问：过山车行情意味着什么")
- education: 学校、考试、招生 (例: "中华女子学院仅1专业招男生")
- science: 技术、网站、宇航 (例: "“手机钱包”亮相科博会")
- society: 社会事件、犯罪、公益 (例: "82岁老太为学生做饭扫地44年")
- politics: 政策、国际关系 (例: "查韦斯称愿为俄罗斯提供空军基地")
- sports: 比赛、运动员、奥运 (例: "卡佩罗：德国脚生猛的原因")
- game: 电子游戏、网游、电竞 (例: "《赤壁OL》攻城战硝烟又起")
- entertainment: 明星、影视、综艺 (例: "冯德伦徐若瑄隔空传情")    """

class BertConfig:
    def __init__(self):
        self.pretrained_model_path = 'model/bert/bert-base-chinese'
        self.batch_size = 2
        self.epochs = 2
        self.padding_size = 32
        self.lr = 5e-5
        self.model_path = 'model/bert/bert_model.bin'
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f'bert_config device={self.device}')
        self.log_batch = 10000


class RfConfig:
    def __init__(self):
        self.preprocessed_train_path = 'data/rf/preprocessed_train.txt'
        self.preprocessed_test_path = 'data/rf/preprocessed_test.txt'
        self.preprocessed_dev_path = 'data/rf/preprocessed_dev.txt'
        self.tfidf_path = 'data/rf/tfidf.pkl'
        self.rf_model_path = 'model/rf/rf_model.pkl'


class FasttextConfig:
    def __init__(self):
        self.train_char_path = 'data/fasttext/train_char.txt'
        self.test_char_path = 'data/fasttext/test_char.txt'
        self.dev_char_path = 'data/fasttext/dev_char.txt'
        self.train_token_path = 'data/fasttext/train_token.txt'
        self.test_token_path = 'data/fasttext/test_token.txt'
        self.dev_token_path = 'data/fasttext/dev_token.txt'

        self.char_model = 'model/fasttext/char_model.bin'
        self.token_model = 'model/fasttext/token_model.bin'
        self.char_auto_model = 'model/fasttext/char_auto_model.bin'
        self.token_auto_model = 'model/fasttext/token_auto_model.bin'


conf = Config()
